package com.lry.basic.algorithm.trie;

import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;

public class Trie {

    private class Node{

        public int use;//字符使用次数
        public int weight;//某个单词的权重
        public TreeMap<Character,Node> next;//映射

        public Node(){
            next = new TreeMap<>();
        }
    }


    private Node root;//根
    private int size;//字典树总单词数


    public Trie(){
        root = new Node();
        size = 0;
    }

    /**
     * 判断某个节点是否是某个单词的结束
     * @param node
     * @return
     */
    private boolean isWord(Node node){
        if(null==node)
            return false;
        return node.weight!=0;
    }

    // 获得Trie中存储的单词数量
    public int size(){
        return size;
    }

    /**
     * 添加一个单词
     * @param word
     */
    public void add(String word){
        char[] cs = checkWord(word);
        add(cs,cs.length);
    }

    /**
     * 添加一个单词
     * @param word 单词
     * @param weight 权重
     */
    public void add(String word, int weight){
        if(0==weight)
            throw new RuntimeException("weight must not zero");
        char[] cs = checkWord(word);
        add(cs,weight);
    }

    private void add(char[] cs, int weight){
        Node cur = root;
        for(char c:cs){
            if(null==cur.next.get(c)){
                cur.next.put(c,new Node());
            }
            cur = cur.next.get(c);
            cur.use++;
        }
        if(!isWord(cur)){
            ++size;
            cur.weight = weight;
        }
    }

    /**
     * 带.的匹配
     * @param pattern
     * @return
     */
    public boolean match(String pattern){
        char[] cs = checkWord(pattern);
        return match(root,cs,0);
    }

    private boolean match(Node node,char[] cs,int index){
        if(index==cs.length){
            return isWord(node);
        }
        char c = cs[index];
        if('.'!=c){
            Node cur = node.next.get(c);
            if(null==cur)
                return false;
            return match(cur,cs,index+1);
        }else{
            for(char ch:node.next.keySet()){
                if(match(node.next.get(ch),cs,index+1)){
                    return true;
                }
            }
            return false;
        }

    }

    /**
     * 是否包含某个单词
     * @param word
     * @return
     */
    public boolean contains(String word){
        char[]cs = checkWord(word);
        Node cur = root;
        for(char c:cs){
            if(null==cur.next.get(c)){
                return false;
            }
            cur = cur.next.get(c);
        }
        return isWord(cur);
    }


    /**
     * 查询是否在Trie中有单词以prefix为前缀
     * @param prefix
     * @return
     */
    public boolean isPrefix(String prefix){
        char[]cs = checkWord(prefix);
        Node cur = root;
        for(char c:cs){
            if(null==cur.next.get(c)){
                return false;
            }
            cur = cur.next.get(c);
        }
        return true;
    }


    private char[] checkWord(String word){
        if(null==word||word.isEmpty())
            throw new RuntimeException("word must not null or empty");
        return word.toCharArray();
    }

    /**
     * 某个单词的权重
     * @param word
     * @return
     */
    public int getWeight(String word){
        char[]cs = checkWord(word);
        Node cur = root;

        for(char c:cs){
            if(cur.next.get(c)==null)
                throw new RuntimeException("could not find word:"+word);
            cur = cur.next.get(c);
        }
        if(!isWord(cur)){
            throw new RuntimeException("could not find word:"+word);
        }
        return cur.weight;
    }


    /**
     * 找出以prefix为前缀的所有单词的权重和
     * @param prefix
     * @return
     */
    public int sumWeight(String prefix){
        char[]cs = checkWord(prefix);
        Node cur = root;
        for(char c:cs){
            if(cur.next.get(c)==null){
                return 0;
            }
            cur = cur.next.get(c);
        }
        return sumWeight(cur);
    }

    private int sumWeight(Node node){
        int ans = node.weight;
        for(char c:node.next.keySet()){
            ans+= sumWeight(node.next.get(c));
        }
        return ans;
    }

    /**
     *
     * @return 所有单词
     */
    public Set<String> getWords(){
        Set<String> words = new TreeSet<String>();
        getWords(root,words,"");
        return words;
    }

    private void getWords(Node node, Set<String> words,String word) {
        if(isWord(node)){
            words.add(word);
        }
        for(char c:node.next.keySet()){
            getWords(node.next.get(c),words,word+c);
        }
    }

    /**
     * 删除某个单词
     * @param word
     * @return
     */
    public boolean delete(String word){
        if(!contains(word)){
            throw new RuntimeException("word is not exists");
        }
        char[] cs = checkWord(word);
        Node cur = root;
        for(char c:cs){
            //如果某个节点只使用了一次，而且只有一个字节点，直接删除即可
            if(cur.next.get(c).use==1&&cur.next.size()==1){
                cur.next.remove(c);
                cur = null;
                break;
            }
            //否则置为下一个节点
            cur = cur.next.get(c);
            cur.use--;
        }
        //如果走完了，是一个单词的话，置为不是单词
        if(isWord(cur)){
            cur.weight = 0;
        }
        size--;
        return true;
    }


}
